#include <ntifs.h>
#define NTSTRSAFE_NO_CB_FUNCTIONS
#include <ntstrsafe.h>
#include "NativeEnums.h"
#include "../HyperPlatform/kernel_stl.h"
#include "main.h"

NTSTATUS NtFileNameToDosFileName(IN PUNICODE_STRING NtFileName, OUT PUNICODE_STRING DosFileName);;

EXTERN_C{

#define MEM_IMAGE 0x1000000 

#define MemorySectionName 2

NTKERNELAPI PVOID PsGetProcessSectionBaseAddress(PEPROCESS Process);

NTKERNELAPI NTSTATUS NTAPI ZwQueryInformationProcess(
	_In_      HANDLE           ProcessHandle,
	_In_      PROCESSINFOCLASS ProcessInformationClass,
	_Out_     PVOID            ProcessInformation,
	_In_      ULONG            ProcessInformationLength,
	_Out_opt_ PULONG           ReturnLength
);

NTSTATUS NTAPI NewZwQueryVirtualMemory(HANDLE ProcessHandle, PVOID BaseAddress, MEMORY_INFORMATION_CLASS_EX MemoryInformationClass, PVOID MemoryInformation, SIZE_T MemoryInformationLength, PSIZE_T ReturnLength);

typedef struct _MEMORY_SECTION_NAME
{
	UNICODE_STRING SectionFileName;
	WCHAR NameBuffer[1];
} MEMORY_SECTION_NAME, *PMEMORY_SECTION_NAME;

BOOLEAN IsProcessIdEqual(HANDLE ProcessId1, HANDLE ProcessId2)
{
	return ((SIZE_T)ProcessId1 / 4 == (SIZE_T)ProcessId2 / 4) ? TRUE : FALSE;
}

NTSTATUS GetCurrentProcessPath(PUNICODE_STRING ProcessName)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	PVOID ImageBase = PsGetProcessSectionBaseAddress(PsGetCurrentProcess());
	PVOID SectionName = ExAllocatePoolWithTag(NonPagedPool, ProcessName->MaximumLength + sizeof(MEMORY_SECTION_NAME), 'TXSB');
	if (!SectionName)
		return status;

	if (ImageBase)
	{
		MEMORY_BASIC_INFORMATION BasicInfo;
		status = NewZwQueryVirtualMemory(NtCurrentProcess(), ImageBase, MemoryBasicInformationEx, &BasicInfo, sizeof(BasicInfo), NULL);
		if (NT_SUCCESS(status) && BasicInfo.Type == MEM_IMAGE)
		{
			status = NewZwQueryVirtualMemory(NtCurrentProcess(), ImageBase, MemoryMappedFilenameInformation, SectionName, ProcessName->MaximumLength + sizeof(MEMORY_SECTION_NAME), NULL);
			if (NT_SUCCESS(status))
			{
				PUNICODE_STRING pSectionFileName = &((PMEMORY_SECTION_NAME)SectionName)->SectionFileName;

				if(!NT_SUCCESS(NtFileNameToDosFileName(pSectionFileName, ProcessName)))
					RtlCopyUnicodeString(ProcessName, pSectionFileName);

				ExFreePoolWithTag(SectionName, 'TXSB');
				return STATUS_SUCCESS;
			}
		}
	}

	ExFreePoolWithTag(SectionName, 'TXSB');
	return status;
}

NTSTATUS GetProcessPathByPID(HANDLE ProcessId, PUNICODE_STRING ProcessName)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	if ((SIZE_T)ProcessId <= 4)
	{
		RtlInitUnicodeString(ProcessName, L"System");
		status = STATUS_SUCCESS;
	}
	else if (IsProcessIdEqual(ProcessId, PsGetCurrentProcessId()))
	{
		status = GetCurrentProcessPath(ProcessName);
	}
	else
	{
		PEPROCESS Process = NULL;
		if (NT_SUCCESS(PsLookupProcessByProcessId(ProcessId, &Process)))
		{
			KAPC_STATE KApc;
			KeStackAttachProcess((PRKPROCESS)Process, &KApc);
			status = GetCurrentProcessPath(ProcessName);
			KeUnstackDetachProcess(&KApc);
			ObDereferenceObject(Process);
		}
	}
	return status;
}

NTSTATUS GetCurrentImagePath(PVOID ImageBase, PUNICODE_STRING ImagePath)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;
	PVOID SectionName = ExAllocatePoolWithTag(NonPagedPool, ImagePath->MaximumLength + sizeof(MEMORY_SECTION_NAME), 'TXSB');
	if (!SectionName)
		return status;

	if (ImageBase)
	{
		MEMORY_BASIC_INFORMATION BasicInfo;
		status = NewZwQueryVirtualMemory(NtCurrentProcess(), ImageBase, MemoryBasicInformationEx, &BasicInfo, sizeof(BasicInfo), NULL);
		if (NT_SUCCESS(status) && BasicInfo.Type == MEM_IMAGE)
		{
			status = NewZwQueryVirtualMemory(NtCurrentProcess(), ImageBase, MemoryMappedFilenameInformation, SectionName, ImagePath->MaximumLength + sizeof(MEMORY_SECTION_NAME), NULL);
			if (NT_SUCCESS(status))
			{
				PUNICODE_STRING pSectionFileName = &((PMEMORY_SECTION_NAME)SectionName)->SectionFileName;

				if (!NT_SUCCESS(NtFileNameToDosFileName(pSectionFileName, ImagePath)))
					RtlCopyUnicodeString(ImagePath, pSectionFileName);

				ExFreePool(SectionName);
				return STATUS_SUCCESS;
			}
		}
	}

	ExFreePool(SectionName);
	return status;
}

NTSTATUS GetImagePathByAddress(HANDLE ProcessId, PVOID ImageBase, PUNICODE_STRING ImagePath)
{
	NTSTATUS status = STATUS_UNSUCCESSFUL;

	if (IsProcessIdEqual(ProcessId, PsGetCurrentProcessId()))
	{
		status = GetCurrentImagePath(ImageBase, ImagePath);
	}
	else
	{
		PEPROCESS Process = NULL;
		if (NT_SUCCESS(PsLookupProcessByProcessId(ProcessId, &Process)))
		{
			KAPC_STATE KApc;
			KeStackAttachProcess((PRKPROCESS)Process, &KApc);
			status = GetCurrentImagePath(ImageBase, ImagePath);
			KeUnstackDetachProcess(&KApc);
			ObDereferenceObject(Process);
		}
	}

	return status;
}

NTSTATUS GetProcessIdByHandle(__in HANDLE ProcessHandle, __out PHANDLE ProcessId)
{
	PROCESS_BASIC_INFORMATION pbi;
	
	NTSTATUS st = ZwQueryInformationProcess(ProcessHandle, ProcessBasicInformation, &pbi, sizeof(pbi), NULL);
	if (NT_SUCCESS(st))
	{
		*ProcessId = (HANDLE)pbi.UniqueProcessId;
		return STATUS_SUCCESS;
	}
	return st;
}

}